#include <opencv2/opencv.hpp>
#include <opencv2/core/core.hpp>
#include <opencv2/highgui/highgui.hpp>
#include <opencv2/imgproc/imgproc.hpp>
#include <iostream>
#include <complex>
#include "cvabs.h"
#include "elementwise.h"

using namespace cv;
using namespace std;

void Dx( const Mat* u, Mat*d );
void Dxt( const Mat* u, Mat* d );
void Dy( Mat* u, Mat* d );
void Dyt( Mat* u, Mat* d );
void shrink2( Mat_<usetype> xs, Mat_<usetype> ys, Mat_<usetype> x, Mat_<usetype> y, float lambda );

typedef complex<float> usetype; // allows easy toggling between float and double

void split_bregman(Mat_<usetype>* R, Mat_<usetype>* f, Mat_<usetype>* u, const float mu, const float lambda, const float gamma, const int nInner, const int nBreg)
{
	const int frows = f->rows;
	const int fcols = f->cols;
	const float scale = sqrt( float(frows * fcols) );
	const Mat f0 = f->clone();

	Mat_<usetype> x(frows, fcols, 0.0);
	Mat_<usetype> y(frows, fcols, 0.0);
	Mat_<usetype> bx(frows, fcols, 0.0);
	Mat_<usetype> by(frows, fcols, 0.0);
	Mat_<usetype> uker(frows, fcols, 0.0);
	(*u) = Mat_<usetype>(frows, fcols, 0.0);
	Mat_<usetype> murf;
	
	// build kernels
	murf = 1.0/scale * mu * elmult(R, f);
	idft(murf, murf);

	uker(0,0) = 4.0;
	uker(0,1) = -1.0;
	uker(1,0) = -1.0;
	uker(frows-1, 0) = -1.0;
	uker(0, fcols-1) = -1.0;
	dft(uker, uker);
	uker = mu * (*R) + lambda * uker;// + gamma;
	Mat ukerplanes[] = {Mat::zeros(frows, fcols, CV_32F), Mat::zeros(frows, fcols, CV_32F)};
	split(uker, ukerplanes);
	ukerplanes[0] = ukerplanes[0] + gamma;
	merge(ukerplanes, 2, uker);

	// do the reconstruction
	Mat_<usetype> rhs, xbxtmp, ybytmp, dxbx, dyby, uf, uabs, dx, dy;
	double minel, maxel;
	for(int outer = 0; outer < nBreg; outer++)
	{
		for(int inner = 0; inner < nInner; inner++)
		{
			// update u
			xbxtmp = x-bx;
			Dxt(&xbxtmp, &dxbx);
			ybytmp = y-by;
			Dyt(&ybytmp, &dyby);

			rhs = murf + lambda * dxbx + lambda * dyby + gamma * (*u);
			dft(rhs, rhs);
			rhs = 1/(scale*scale) * eldiv(&rhs, &uker);
			idft(rhs, *u);	
			//cvabs(u, &uabs);
			//minMaxIdx(uabs, &minel, &maxel);

			// update x and y
			Dx(u, &dx);
			Dy(u, &dy);
			//dbx = dx + bx;
			//dby = dy + by;
			shrink2(x, y, dx + bx, dy + by, 1.0/lambda);
			bx += dx - x;
			by += dy - y;
		}
		dft(*u, uf);
		(*f) = (*f) + f0 - 1./scale * elmult(R, &uf);
		idft((mu / scale) * elmult(R, f), murf);
	}
}

void Dx( const Mat *u, Mat *d )
{
	const int urows = u->rows;
	const int ucols = u->cols;

	*d = Mat_<usetype>(urows, ucols, 0.0);

	d->col(0) = u->col(0) - u->col(ucols-1);
	(*d)(Range::all(), Range(1, ucols)) = (*u)(Range::all(), Range(1, ucols)) - (*u)(Range::all(), Range(0, ucols-1));
}

void Dxt( const Mat* u, Mat* d )
{
	const int urows = u->rows;
	const int ucols = u->cols;

	*d = Mat_<usetype>(urows, ucols, 0.0);

	(*d)(Range::all(), Range(0, ucols-1)) = (*u)(Range::all(), Range(0, ucols-1)) - (*u)(Range::all(), Range(1, ucols));
	d->col(ucols-1) = u->col(ucols-1) - u->col(0);
}

void Dy( Mat* u, Mat* d )
{
	const int urows = u->rows;
	const int ucols = u->cols;

	*d = Mat_<usetype>(urows, ucols, 0.0);

	d->row(0) = u->row(0) - u->row(urows-1);
	(*d)(Range(1, urows), Range::all()) = (*u)(Range(1, urows), Range::all()) - (*u)(Range(0, urows-1), Range::all());
}

void Dyt( Mat* u, Mat* d )
{
	const int urows = u->rows;
	const int ucols = u->cols;
	
	*d = Mat_<usetype>(urows, ucols, 0.0);

	(*d)(Range(0, urows-1), Range::all()) = (*u)(Range(0, urows-1), Range::all()) - (*u)(Range(1, urows), Range::all());
	d->row(urows-1) = u->row(urows-1) - u->row(0);
}


void shrink2( Mat_<usetype> xs, Mat_<usetype> ys, Mat_<usetype> x, Mat_<usetype> y, float lambda )
{
	Mat xabs, yabs, s, ss, sless;
	Mat_<usetype>  s2, sss, spos;
	cvabs(&x, &xabs);
	pow(xabs, 2, xabs);
	cvabs(&y, &yabs);
	pow(yabs, 2, yabs);
	pow(xabs + yabs, 0.5, s);

	ss = s-lambda;
	Mat tmp;
	Mat_<usetype> tmp2;
	threshold(ss, tmp, 0, 1, THRESH_TOZERO);
	/*compare(ss, 0, spos, CMP_GT);
	spos.convertTo(spos, CV_32FC2);
	ss = ss.mul(spos/255.0);*/
	threshold(s, sless, lambda, 1, THRESH_BINARY_INV);
	/*compare(s, lambda, sless, CMP_LT);
	sless.convertTo(sless, CV_32FC2);
	sless /= 255.0;*/
	s = s + sless;
	merge2(s, Mat::zeros(s.size(), CV_32F), s2);
	merge2(tmp, Mat::zeros(s.size(), CV_32F), tmp2);
	sss = eldiv(&tmp2,&s2);
	xs = elmult(&sss, &x);
	ys = elmult(&sss, &y);
}